import re

import torch
import numpy as np

ff = open("housing.data", encoding="utf-8").readlines()
data = []
for item in ff:
    out = re.sub(r"\s{2,}", " ", item).strip()
    data.append(out.split(" "))
data = np.array(data).astype(np.float64)

Y = data[:, -1]
X = data[:, 0:-1]

X_train = X[0:496, ...]
Y_train = Y[0:496, ...]
X_test = X[496:, ...]
Y_test = Y[496:, ...]


class Net(torch.nn.Module):
    def __init__(self, n_feature, n_output):
        super(Net, self).__init__()
        self.hidden = torch.nn.Linear(n_feature, 100)
        self.predict = torch.nn.Linear(100, n_output)

    def forward(self, x):
        out = self.hidden(x)
        out = torch.relu(out)
        out = self.predict(out)
        return out


net = Net(13, 1)

loss_func = torch.nn.MSELoss()

optimizer = torch.optim.SGD(net.parameters(), lr=0.001)

for i in range(10000):
    x_data = torch.tensor(X_train, dtype=torch.float32)
    y_data = torch.tensor(Y_train, dtype=torch.float32)
    pred = net.forward(x_data).squeeze()
    loss = loss_func(pred, y_data)
    optimizer.zero_grad()
    loss.backward()
    optimizer.step()
    print(loss)

    print(loss_func(net.forward(torch.tensor(X_test, dtype=torch.float32)).squeeze(), torch.tensor(Y_test, dtype=torch.float32)))
    print("*"*100)

torch.save(net,"model.pkl") # 这种保存文件大，但是读取出来就能用
torch.save(net.state_dict(),"op.pkl") # 这种只保存数据，用的时候先定义模型再加载数据